import os
import warnings
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
import os

# 设置绘图风格
sns.set(style="whitegrid", palette="pastel")
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 创建分析结果目录结构
analysis_dir = "analysis_results"
os.makedirs(analysis_dir, exist_ok=True)

# 为每种疾病创建子目录
stroke_dir = os.path.join(analysis_dir, "stroke")
heart_dir = os.path.join(analysis_dir, "heart")
cirrhosis_dir = os.path.join(analysis_dir, "cirrhosis")
comparison_dir = os.path.join(analysis_dir, "comparison")

for dir_path in [stroke_dir, heart_dir, cirrhosis_dir, comparison_dir]:
    os.makedirs(dir_path, exist_ok=True)

# 加载处理后的数据
print("加载数据集...")
stroke_df = pd.read_csv('output/stroke_processed.csv')
heart_df = pd.read_csv('output/heart_processed.csv')
cirrhosis_df = pd.read_csv('output/cirrhosis_processed.csv')

# 打印数据集列名以进行调试
print("\n数据集列名:")
print("中风数据集:", stroke_df.columns.tolist())
print("心脏病数据集:", heart_df.columns.tolist())
print("肝硬化数据集:", cirrhosis_df.columns.tolist())


# 识别目标列
def find_target_column(df, dataset_name):
    """识别目标列"""
    possible_targets = {
        'stroke': ['stroke', 'target', 'disease', 'outcome'],
        'heart': ['target', 'disease', 'condition', 'heart_disease', 'outcome'],
        'cirrhosis': ['Stage', 'stage', 'disease_stage', 'status', 'outcome']
    }

    # 尝试识别目标列
    for name in possible_targets[dataset_name]:
        if name in df.columns:
            print(f"在 {dataset_name} 数据集中找到目标列: {name}")
            return name

    # 如果找不到，尝试最后一列
    last_col = df.columns[-1]
    print(f"警告: 在 {dataset_name} 数据集中未找到目标列，使用最后一列 '{last_col}' 作为目标列")
    return last_col


# 为每个数据集识别目标列
stroke_target = find_target_column(stroke_df, 'stroke')
heart_target = find_target_column(heart_df, 'heart')
cirrhosis_target = find_target_column(cirrhosis_df, 'cirrhosis')

# 删除目标列有缺失值的行，防止 astype(int) 报错
stroke_df = stroke_df[stroke_df[stroke_target].notna()]
heart_df = heart_df[heart_df[heart_target].notna()]
cirrhosis_df = cirrhosis_df[cirrhosis_df[cirrhosis_target].notna()]

# 确保目标列的数据类型正确
stroke_df[stroke_target] = stroke_df[stroke_target].astype(int)
heart_df[heart_target] = heart_df[heart_target].astype(int)
cirrhosis_df[cirrhosis_target] = cirrhosis_df[cirrhosis_target].astype(int)


def save_fig(fig, filename, directory=None):
    """保存图表到文件"""
    if directory is None:
        directory = analysis_dir
    fig.savefig(os.path.join(directory, filename), bbox_inches='tight', dpi=300)
    plt.close(fig)


def analyze_stroke():
    """分析中风数据集"""
    print("\n" + "=" * 50)
    print("中风数据集分析")
    print("=" * 50)

    # 使用正确的目标列名
    target_col = stroke_target

    # 1. 基本统计分析
    print("\n[基本统计]")
    print(stroke_df.describe().round(2))

    # 2. 中风分布
    stroke_counts = stroke_df[target_col].value_counts()
    fig, ax = plt.subplots(figsize=(8, 6))
    sns.barplot(x=stroke_counts.index, y=stroke_counts.values, hue=stroke_counts.index,
                palette="viridis", legend=False)
    plt.title('中风分布 (0=未中风, 1=中风)')
    plt.ylabel('人数')
    save_fig(fig, 'stroke_distribution.png', stroke_dir)

    # 3. 年龄分布与中风关系
    fig, ax = plt.subplots(figsize=(10, 6))
    sns.histplot(data=stroke_df, x='age', hue=target_col, bins=30, kde=True,
                 palette={0: "skyblue", 1: "coral"}, element="step", alpha=0.7)
    plt.title('年龄分布与中风关系')
    plt.xlabel('年龄')
    save_fig(fig, 'stroke_age_distribution.png', stroke_dir)

    # 4. BMI与中风关系
    fig, ax = plt.subplots(1, 2, figsize=(14, 6))
    sns.boxplot(x=target_col, y='bmi', data=stroke_df, hue=target_col,
                palette={0: "skyblue", 1: "coral"}, legend=False, ax=ax[0])
    ax[0].set_title('BMI与中风关系')
    ax[0].set_xlabel('中风 (0=否, 1=是)')
    ax[0].set_ylabel('BMI')

    sns.scatterplot(x='age', y='bmi', hue=target_col, data=stroke_df,
                    palette={0: "skyblue", 1: "red"}, alpha=0.6, ax=ax[1])
    ax[1].set_title('年龄-BMI与中风关系')
    save_fig(fig, 'stroke_bmi_relationship.png', stroke_dir)

    # 5. 性别与中风关系
    if 'gender' in stroke_df.columns:
        gender_stroke = stroke_df.groupby(['gender', target_col]).size().unstack()
        gender_stroke_percent = gender_stroke.div(gender_stroke.sum(axis=1), axis=0) * 100

        fig, ax = plt.subplots(figsize=(10, 6))
        gender_stroke_percent.plot(kind='bar', stacked=True, color=['skyblue', 'coral'], ax=ax)
        plt.title('性别与中风关系')
        plt.ylabel('百分比 (%)')
        plt.xlabel('性别')
        plt.legend(['未中风', '中风'], title='中风状态')
        save_fig(fig, 'stroke_gender_relationship.png', stroke_dir)
    else:
        print("警告: 中风数据集缺少'gender'列")

    # 6. 吸烟状态与中风关系
    if 'smoking_status' in stroke_df.columns:
        fig, ax = plt.subplots(figsize=(10, 6))
        sns.countplot(x='smoking_status', hue=target_col, data=stroke_df,
                      palette={0: "skyblue", 1: "coral"})
        plt.title('吸烟状态与中风关系')
        plt.xlabel('吸烟状态')
        plt.ylabel('人数')
        plt.legend(['未中风', '中风'], title='中风状态')
        save_fig(fig, 'stroke_smoking_relationship.png', stroke_dir)
    else:
        print("警告: 中风数据集缺少'smoking_status'列")

    # 7. 血糖水平与中风关系
    if 'avg_glucose_level' in stroke_df.columns:
        fig, ax = plt.subplots(figsize=(10, 6))
        sns.violinplot(x=target_col, y='avg_glucose_level', data=stroke_df, hue=target_col,
                       palette={0: "skyblue", 1: "coral"}, inner="quartile", legend=False)
        plt.title('平均血糖水平与中风关系')
        plt.xlabel('中风 (0=否, 1=是)')
        plt.ylabel('平均血糖水平')
        save_fig(fig, 'stroke_glucose_relationship.png', stroke_dir)
    else:
        print("警告: 中风数据集缺少'avg_glucose_level'列")

    # 8. 特征相关性分析
    numeric_cols = stroke_df.select_dtypes(include=np.number).columns
    corr = stroke_df[numeric_cols].corr()

    fig, ax = plt.subplots(figsize=(12, 10))

    # 修复热力图问题：移除mask参数，使用完整矩阵
    sns.heatmap(corr, annot=True, fmt=".2f", cmap="coolwarm",
                cbar_kws={"shrink": 0.8}, linewidths=0.5)
    plt.title('中风数据集特征相关性热力图')
    save_fig(fig, 'stroke_correlation.png', stroke_dir)

    # 9. 多因素回归分析
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import LabelEncoder

    # 准备数据
    stroke_df_encoded = stroke_df.copy()
    le = LabelEncoder()

    # 识别分类列
    categorical_cols = stroke_df_encoded.select_dtypes(include=['object', 'category']).columns
    for col in categorical_cols:
        stroke_df_encoded[col] = le.fit_transform(stroke_df_encoded[col].astype(str))

    # 移除ID列和目标列
    features = [col for col in stroke_df_encoded.columns if col not in ['id', target_col]]
    X = stroke_df_encoded[features]
    y = stroke_df_encoded[target_col]

    # 移除含有NaN的样本，防止LogisticRegression报错
    X = X.dropna()
    y = y.loc[X.index]

    # 训练逻辑回归模型
    model = LogisticRegression(max_iter=1000)
    model.fit(X, y)

    # 获取特征重要性
    importance = pd.DataFrame({
        'Feature': X.columns,
        'Importance': model.coef_[0]
    }).sort_values('Importance', ascending=False)

    fig, ax = plt.subplots(figsize=(12, 8))
    sns.barplot(x='Importance', y='Feature', data=importance, palette="viridis", legend=False)
    plt.title('影响中风概率的特征重要性')
    plt.xlabel('重要性 (回归系数)')
    plt.ylabel('特征')
    save_fig(fig, 'stroke_feature_importance.png', stroke_dir)

    # 返回分析结果
    return {
        'stroke_rate': stroke_counts[1] / len(stroke_df) if 1 in stroke_counts else 0,
        'mean_age': stroke_df['age'].mean() if 'age' in stroke_df.columns else None,
        'mean_bmi': stroke_df['bmi'].mean() if 'bmi' in stroke_df.columns else None,
        'feature_importance': importance
    }


def analyze_heart():
    """分析心脏病数据集"""
    print("\n" + "=" * 50)
    print("心脏病数据集分析")
    print("=" * 50)

    # 使用正确的目标列名
    target_col = heart_target

    # 1. 基本统计分析
    print("\n[基本统计]")
    print(heart_df.describe().round(2))

    # 2. 心脏病分布
    heart_counts = heart_df[target_col].value_counts()
    fig, ax = plt.subplots(figsize=(8, 6))
    sns.barplot(x=heart_counts.index, y=heart_counts.values, hue=heart_counts.index,
                palette="magma", legend=False)
    plt.title('心脏病分布 (0=无心脏病, 1=有心脏病)')
    plt.ylabel('人数')
    save_fig(fig, 'heart_distribution.png', heart_dir)

    # 3. 年龄分布与心脏病关系
    if 'Age' in heart_df.columns:
        fig, ax = plt.subplots(figsize=(10, 6))
        sns.histplot(data=heart_df, x='Age', hue=target_col, bins=30, kde=True,
                     palette={0: "skyblue", 1: "purple"}, element="step", alpha=0.7)
        plt.title('年龄分布与心脏病关系')
        plt.xlabel('年龄')
        save_fig(fig, 'heart_age_distribution.png', heart_dir)
    else:
        print("警告: 心脏病数据集缺少'Age'列")

    # 4. 胆固醇与心脏病关系
    if 'Cholesterol' in heart_df.columns and 'Age' in heart_df.columns:
        fig, ax = plt.subplots(1, 2, figsize=(14, 6))
        sns.boxplot(x=target_col, y='Cholesterol', data=heart_df, hue=target_col,
                    palette={0: "skyblue", 1: "purple"}, legend=False, ax=ax[0])
        ax[0].set_title('胆固醇水平与心脏病关系')
        ax[0].set_xlabel('心脏病 (0=无, 1=有)')
        ax[0].set_ylabel('胆固醇')

        sns.scatterplot(x='Age', y='Cholesterol', hue=target_col, data=heart_df,
                        palette={0: "skyblue", 1: "purple"}, alpha=0.6, ax=ax[1])
        ax[1].set_title('年龄-胆固醇与心脏病关系')
        save_fig(fig, 'heart_cholesterol_relationship.png', heart_dir)
    else:
        print("警告: 心脏病数据集缺少'Cholesterol'或'Age'列")

    # 5. 性别与心脏病关系
    if 'Sex' in heart_df.columns:
        # 确保性别列是分类类型
        heart_df['Sex'] = heart_df['Sex'].astype(str)

        gender_heart = heart_df.groupby(['Sex', target_col]).size().unstack()
        gender_heart_percent = gender_heart.div(gender_heart.sum(axis=1), axis=0) * 100

        fig, ax = plt.subplots(figsize=(10, 6))
        gender_heart_percent.plot(kind='bar', stacked=True, color=['skyblue', 'purple'], ax=ax)
        plt.title('性别与心脏病关系')
        plt.ylabel('百分比 (%)')
        plt.xlabel('性别')
        plt.legend(['无心脏病', '有心脏病'], title='心脏病状态')
        save_fig(fig, 'heart_gender_relationship.png', heart_dir)
    else:
        print("警告: 心脏病数据集缺少'Sex'列")

    # 6. 胸痛类型与心脏病关系
    if 'ChestPainType' in heart_df.columns:
        # 确保胸痛类型是分类类型
        heart_df['ChestPainType'] = heart_df['ChestPainType'].astype(str)

        fig, ax = plt.subplots(figsize=(10, 6))
        sns.countplot(x='ChestPainType', hue=target_col, data=heart_df,
                      palette={0: "skyblue", 1: "purple"})
        plt.title('胸痛类型与心脏病关系')
        plt.xlabel('胸痛类型')
        plt.ylabel('人数')
        plt.legend(['无心脏病', '有心脏病'], title='心脏病状态')
        save_fig(fig, 'heart_chest_pain_relationship.png', heart_dir)
    else:
        print("警告: 心脏病数据集缺少'ChestPainType'列")

    # 7. 特征相关性分析
    numeric_cols = heart_df.select_dtypes(include=np.number).columns
    corr = heart_df[numeric_cols].corr()

    fig, ax = plt.subplots(figsize=(12, 10))

    # 修复热力图问题：移除mask参数，使用完整矩阵
    sns.heatmap(corr, annot=True, fmt=".2f", cmap="coolwarm",
                cbar_kws={"shrink": 0.8}, linewidths=0.5)
    plt.title('心脏病数据集特征相关性热力图')
    save_fig(fig, 'heart_correlation.png', heart_dir)

    # 8. 多因素回归分析
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import LabelEncoder

    # 创建编码版本
    heart_df_encoded = heart_df.copy()

    # 对分类列进行编码
    categorical_cols = heart_df.select_dtypes(include=['object', 'category']).columns
    for col in categorical_cols:
        heart_df_encoded[col] = LabelEncoder().fit_transform(heart_df_encoded[col].astype(str))

    # 移除目标列
    features = [col for col in heart_df_encoded.columns if col != target_col]
    X = heart_df_encoded[features]
    y = heart_df_encoded[target_col]

    # 移除含有NaN的样本，防止LogisticRegression报错
    X = X.dropna()
    y = y.loc[X.index]

    model = LogisticRegression(max_iter=1000)
    model.fit(X, y)

    importance = pd.DataFrame({
        'Feature': X.columns,
        'Importance': model.coef_[0]
    }).sort_values('Importance', ascending=False)

    fig, ax = plt.subplots(figsize=(12, 8))
    sns.barplot(x='Importance', y='Feature', data=importance, palette="magma", legend=False)
    plt.title('影响心脏病概率的特征重要性')
    plt.xlabel('重要性 (回归系数)')
    plt.ylabel('特征')
    save_fig(fig, 'heart_feature_importance.png', heart_dir)

    # 返回分析结果
    return {
        'heart_disease_rate': heart_counts[1] / len(heart_df) if 1 in heart_counts else 0,
        'mean_age': heart_df['Age'].mean() if 'Age' in heart_df.columns else None,
        'mean_chol': heart_df['Cholesterol'].mean() if 'Cholesterol' in heart_df.columns else None,
        'feature_importance': importance
    }


def analyze_cirrhosis():
    """分析肝硬化数据集"""
    print("\n" + "=" * 50)
    print("肝硬化数据集分析")
    print("=" * 50)

    # 使用正确的目标列名
    target_col = cirrhosis_target

    # 1. 基本统计分析
    print("\n[基本统计]")
    print(cirrhosis_df.describe().round(2))

    # 2. 肝硬化阶段分布
    fig, ax = plt.subplots(figsize=(10, 6))
    sns.countplot(x=target_col, data=cirrhosis_df, hue=target_col, palette="rocket", legend=False)
    plt.title('肝硬化阶段分布')
    plt.xlabel('阶段')
    plt.ylabel('人数')
    save_fig(fig, 'cirrhosis_stage_distribution.png', cirrhosis_dir)

    # 3. 年龄分布与肝硬化阶段关系
    if 'Age' in cirrhosis_df.columns:
        fig, ax = plt.subplots(figsize=(12, 7))
        sns.boxplot(x=target_col, y='Age', data=cirrhosis_df, hue=target_col, palette="rocket", legend=False)
        plt.title('年龄与肝硬化阶段关系')
        plt.xlabel('肝硬化阶段')
        plt.ylabel('年龄')
        save_fig(fig, 'cirrhosis_age_stage.png', cirrhosis_dir)
    else:
        print("警告: 肝硬化数据集缺少'Age'列")

    # 4. 关键医学指标与肝硬化阶段关系
    fig, ax = plt.subplots(2, 2, figsize=(16, 12))
    # 胆红素
    if 'Bilirubin' in cirrhosis_df.columns:
        sns.boxplot(x=target_col, y='Bilirubin', data=cirrhosis_df, hue=target_col, palette="rocket",
                    legend=False, ax=ax[0, 0])
        ax[0, 0].set_title('胆红素水平与肝硬化阶段')
    # 白蛋白
    if 'Albumin' in cirrhosis_df.columns:
        sns.boxplot(x=target_col, y='Albumin', data=cirrhosis_df, hue=target_col, palette="rocket",
                    legend=False, ax=ax[0, 1])
        ax[0, 1].set_title('白蛋白水平与肝硬化阶段')
    # 凝血酶原时间
    if 'Prothrombin' in cirrhosis_df.columns:
        sns.boxplot(x=target_col, y='Prothrombin', data=cirrhosis_df, hue=target_col, palette="rocket",
                    legend=False, ax=ax[1, 0])
        ax[1, 0].set_title('凝血酶原时间与肝硬化阶段')
    # 铜含量
    if 'Copper' in cirrhosis_df.columns:
        sns.boxplot(x=target_col, y='Copper', data=cirrhosis_df, hue=target_col, palette="rocket",
                    legend=False, ax=ax[1, 1])
        ax[1, 1].set_title('铜含量与肝硬化阶段')
    save_fig(fig, 'cirrhosis_medical_indicators.png', cirrhosis_dir)

    # 5. 药物与肝硬化阶段关系
    if 'Drug' in cirrhosis_df.columns:
        # 确保药物类型是字符串
        cirrhosis_df['Drug'] = cirrhosis_df['Drug'].astype(str)

        fig, ax = plt.subplots(figsize=(12, 7))
        sns.countplot(x='Drug', hue=target_col, data=cirrhosis_df, palette="rocket")
        plt.title('药物类型与肝硬化阶段关系')
        plt.xlabel('药物类型')
        plt.ylabel('人数')
        plt.legend(title='肝硬化阶段', bbox_to_anchor=(1.05, 1), loc='upper left')
        save_fig(fig, 'cirrhosis_drug_stage.png', cirrhosis_dir)
    else:
        print("警告: 肝硬化数据集缺少'Drug'列")

    # 6. 性别与肝硬化关系
    if 'Sex' in cirrhosis_df.columns:
        # 确保性别是字符串
        cirrhosis_df['Sex'] = cirrhosis_df['Sex'].astype(str)

        gender_cirrhosis = cirrhosis_df.groupby(['Sex', target_col]).size().unstack()
        gender_cirrhosis_percent = gender_cirrhosis.div(gender_cirrhosis.sum(axis=1), axis=0) * 100

        fig, ax = plt.subplots(figsize=(12, 7))
        gender_cirrhosis_percent.plot(kind='bar', stacked=True, cmap="rocket", ax=ax)
        plt.title('性别与肝硬化阶段关系')
        plt.ylabel('百分比 (%)')
        plt.xlabel('性别')
        plt.legend(title='肝硬化阶段', bbox_to_anchor=(1.05, 1), loc='upper left')
        save_fig(fig, 'cirrhosis_gender_stage.png', cirrhosis_dir)
    else:
        print("警告: 肝硬化数据集缺少'Sex'列")

    # 7. 特征相关性分析
    numeric_cols = cirrhosis_df.select_dtypes(include=np.number).columns
    corr = cirrhosis_df[numeric_cols].corr()

    fig, ax = plt.subplots(figsize=(14, 12))

    # 修复热力图问题：移除mask参数，使用完整矩阵
    sns.heatmap(corr, annot=True, fmt=".2f", cmap="coolwarm",
                cbar_kws={"shrink": 0.8}, linewidths=0.5)
    plt.title('肝硬化数据集特征相关性热力图')
    save_fig(fig, 'cirrhosis_correlation.png', cirrhosis_dir)

    # 8. 多因素回归分析
    from sklearn.linear_model import LinearRegression
    from sklearn.preprocessing import LabelEncoder

    cirrhosis_df_encoded = cirrhosis_df.copy()
    le = LabelEncoder()

    # 对分类列进行编码
    categorical_cols = cirrhosis_df.select_dtypes(include=['object', 'category']).columns
    for col in categorical_cols:
        cirrhosis_df_encoded[col] = le.fit_transform(cirrhosis_df_encoded[col].astype(str))

    # 移除ID列和目标列
    features = [col for col in cirrhosis_df_encoded.columns if col not in ['ID', target_col]]
    X = cirrhosis_df_encoded[features]
    y = cirrhosis_df_encoded[target_col]

    # 移除含有NaN的样本，防止LinearRegression报错
    X = X.dropna()
    y = y.loc[X.index]

    model = LinearRegression()
    model.fit(X, y)

    importance = pd.DataFrame({
        'Feature': X.columns,
        'Importance': model.coef_
    }).sort_values('Importance', ascending=False)

    fig, ax = plt.subplots(figsize=(14, 10))
    sns.barplot(x='Importance', y='Feature', data=importance, palette="rocket", legend=False)
    plt.title('影响肝硬化阶段的因素重要性')
    plt.xlabel('重要性 (回归系数)')
    plt.ylabel('特征')
    save_fig(fig, 'cirrhosis_feature_importance.png', cirrhosis_dir)

    # 返回分析结果
    return {
        'mean_stage': cirrhosis_df[target_col].mean(),
        'mean_age': cirrhosis_df['Age'].mean() if 'Age' in cirrhosis_df.columns else None,
        'mean_bilirubin': cirrhosis_df['Bilirubin'].mean() if 'Bilirubin' in cirrhosis_df.columns else None,
        'feature_importance': importance
    }


def compare_disease_factors(stroke_results, heart_results, cirrhosis_results):
    """比较三种疾病的影响因素"""
    # 创建对比数据
    diseases = ['中风', '心脏病', '肝硬化']
    disease_rates = [
        stroke_results.get('stroke_rate', 0),
        heart_results.get('heart_disease_rate', 0),
        cirrhosis_results.get('mean_stage', 0) / 4  # 标准化为0-1范围
    ]

    mean_ages = [
        stroke_results.get('mean_age', 0),
        heart_results.get('mean_age', 0),
        cirrhosis_results.get('mean_age', 0)
    ]

    # 绘制疾病患病率比较
    fig, ax = plt.subplots(figsize=(10, 6))
    sns.barplot(x=diseases, y=disease_rates, palette="Set2", legend=False)
    plt.title('三种疾病患病率比较')
    plt.ylabel('患病率/严重程度')
    plt.ylim(0, 1)
    save_fig(fig, 'disease_prevalence_comparison.png', comparison_dir)

    # 绘制平均年龄比较
    fig, ax = plt.subplots(figsize=(10, 6))
    sns.barplot(x=diseases, y=mean_ages, palette="Set2", legend=False)
    plt.title('三种疾病患者平均年龄比较')
    plt.ylabel('平均年龄')
    save_fig(fig, 'disease_mean_age_comparison.png', comparison_dir)

    # 创建特征重要性对比数据
    top_features = {}
    for disease, results in zip(diseases, [stroke_results, heart_results, cirrhosis_results]):
        if 'feature_importance' in results:
            # 取绝对值并按降序排序
            results['feature_importance']['Abs_Importance'] = abs(results['feature_importance']['Importance'])
            top_features[disease] = results['feature_importance'].sort_values('Abs_Importance', ascending=False).head(5)

    # 创建合并的特征重要性数据框
    comparison_df = pd.DataFrame()
    for disease, df in top_features.items():
        temp_df = df[['Feature', 'Importance']].copy()
        temp_df['Disease'] = disease
        comparison_df = pd.concat([comparison_df, temp_df])

    # 绘制特征重要性对比
    if not comparison_df.empty:
        fig, ax = plt.subplots(figsize=(14, 10))
        sns.barplot(x='Importance', y='Feature', hue='Disease', data=comparison_df, palette="Set2")
        plt.title('三种疾病前五位影响因素比较')
        plt.ylabel('特征')
        plt.xlabel('特征重要性 (回归系数)')
        plt.legend(title='疾病类型')
        save_fig(fig, 'disease_feature_comparison.png', comparison_dir)
    else:
        print("警告: 无法创建特征重要性比较图")


# 执行分析
try:
    print("\n开始分析中风数据集...")
    stroke_results = analyze_stroke()

    print("\n开始分析心脏病数据集...")
    heart_results = analyze_heart()

    print("\n开始分析肝硬化数据集...")
    cirrhosis_results = analyze_cirrhosis()

    print("\n比较三种疾病的影响因素...")
    compare_disease_factors(stroke_results, heart_results, cirrhosis_results)

    print("\n分析完成！所有图表已保存至 'analysis_results' 目录")
except Exception as e:
    print(f"分析过程中发生错误: {e}")
    import traceback

    traceback.print_exc()